import torch

if torch.cuda.is_available():
    print("CUDA is available. GPU count:", torch.cuda.device_count())
else:
    print("CUDA is not available.")

